from sqlalchemy import create_engine, MetaData, Table
from sqlalchemy.ext.automap import automap_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy import inspect

# 替换为你的数据库连接字符串
DATABASE_URI = 'mysql+pymysql://root:root%40buaa@172.1.1.82/parking-dev'

# 指定要生成模型的表名列表
TABLES_TO_GENERATE = ['sys_user_post']  # 将 'table1' 和 'table2' 替换为你想要的表名
# 定义输出文件
output_file = './model/entity/sys_user_post.py'

# 创建引擎
engine = create_engine(DATABASE_URI)

# 反射已有的数据库表格
Base = automap_base()
Base.prepare(autoload_with=engine)

# 打印所有映射的表名（可选）
for table_name in Base.classes.keys():
    print(table_name)



with open(output_file, 'w') as f:
    # 写入必要的导入语句
    f.write('from sqlalchemy import BIGINT, CHAR, DATETIME, VARCHAR,Column, Integer, String, ForeignKey, DateTime, Float, Text, Boolean\n')
    f.write('from sqlalchemy.ext.declarative import declarative_base\n\n')
    f.write('Base = declarative_base()\n\n')

    for table_name in TABLES_TO_GENERATE:
        if table_name in Base.classes:
            table_class = Base.classes[table_name]
            mapper = inspect(table_class)
            columns = [column.key for column in mapper.columns]

            f.write(f'class {table_name}(Base):\n')
            f.write('    __tablename__ = "{}"\n'.format(table_name))

            for column in columns:
                col = getattr(mapper.columns, column)
                type_name = type(col.type).__name__

                # 简单处理一些常见的数据类型
                args = []
                if col.primary_key:
                    args.append('primary_key=True')
                if not col.nullable:
                    args.append('nullable=False')
                if col.default is not None:
                    args.append(f'default={col.default.arg!r}')
                if col.foreign_keys:
                    for fk in col.foreign_keys:
                        args.append(f'ForeignKey("{fk.target_fullname}")')

                arg_str = ', '.join(args)
                f.write('    {} = Column({}, {})\n'.format(
                    col.name,
                    type_name,
                    arg_str
                ))
            f.write('\n')
        else:
            print(f"Warning: Table '{table_name}' does not exist in the database.")

print(f"Models have been written to {output_file}")